#include <immintrin.h>
#include <stdio.h>
inline void avx2_add_eltwise(const float *arr1, const float *arr2, float *result, int len)
{
    // 主循环：每次处理 8 个 float（AVX2 的 256 位寄存器）
    int i = 0;
    // 处理剩余元素（不足 8 个的部分）
    for (; i < len; ++i)
    {
        result[i] = arr1[i] + arr2[i];
    }
}

// 函数：标量 float 与 float 数组逐元素相加（AVX 优化）
void avx2_add_scalar_vector(float scalar, const float *arr, float *result, int len)
{
    // 主循环：每次处理 8 个 float（AVX 的 256 位寄存器）
    int i = 0;

    // 处理剩余元素（不足 8 个的部分）
    for (; i < len; ++i)
    {
        result[i] = scalar + arr[i];
    }
}

#define INNER_LOOP_ELTWISE                                                                   \
    avx2_add_eltwise(a_data + a_idx_part0 + a_idx_part1, b_data + b_idx_part0 + b_idx_part1, \
                     c_data + c_idx_part0 + c_idx_part1, c_shape[2]);
#define INNER_LOOP_BROADCAST_A2B                                                                      \
    avx2_add_scalar_vector(*(a_data + a_idx_part0 + a_idx_part1), b_data + b_idx_part0 + b_idx_part1, \
                           c_data + c_idx_part0 + c_idx_part1, c_shape[2]);

#define INNER_LOOP_BROADCAST_B2A                                                                      \
    avx2_add_scalar_vector(*(b_data + b_idx_part0 + b_idx_part1), a_data + a_idx_part0 + a_idx_part1, \
                           c_data + c_idx_part0 + c_idx_part1, c_shape[2]);

#define INNER_LOOP_UNIVERSEL                                       \
    for (int k = 0; k < c_shape[2]; k++)                           \
    {                                                              \
        const int a_idx_part2 = k * adjusted_a_stride[2];          \
        const int b_idx_part2 = k * adjusted_b_stride[2];          \
        const int c_idx_part2 = k * c_stride[2];                   \
        const int a_idx = a_idx_part0 + a_idx_part1 + a_idx_part2; \
        const int b_idx = b_idx_part0 + b_idx_part1 + b_idx_part2; \
        const int c_idx = c_idx_part0 + c_idx_part1 + c_idx_part2; \
        c_data[c_idx] = a_data[a_idx] + b_data[b_idx];             \
    }

#define BROADCAST_ADD_FUNC(suffix, INNER_LOOP)                                  \
    void broadcast_add_##suffix(float *a_data, int a_shape[3], int a_stride[3], \
                                float *b_data, int b_shape[3], int b_stride[3], \
                                float *c_data, int c_shape[3], int c_stride[3]) \
    {                                                                           \
        /* 调整后的步长计算 */                                                  \
        int adjusted_a_stride[3], adjusted_b_stride[3];                         \
        for (int d = 0; d < 3; d++)                                             \
        {                                                                       \
            adjusted_a_stride[d] = (a_shape[d] == 1) ? 0 : a_stride[d];         \
            adjusted_b_stride[d] = (b_shape[d] == 1) ? 0 : b_stride[d];         \
        }                                                                       \
        /* 遍历输出空间 */                                                      \
        for (int i = 0; i < c_shape[0]; i++)                                    \
        {                                                                       \
            const int a_idx_part0 = i * adjusted_a_stride[0];                   \
            const int b_idx_part0 = i * adjusted_b_stride[0];                   \
            const int c_idx_part0 = i * c_stride[0];                            \
            for (int j = 0; j < c_shape[1]; j++)                                \
            {                                                                   \
                const int a_idx_part1 = j * adjusted_a_stride[1];               \
                const int b_idx_part1 = j * adjusted_b_stride[1];               \
                const int c_idx_part1 = j * c_stride[1];                        \
                INNER_LOOP                                                      \
            }                                                                   \
        }                                                                       \
    }

BROADCAST_ADD_FUNC(universal, INNER_LOOP_UNIVERSEL)
BROADCAST_ADD_FUNC(eltwise, INNER_LOOP_ELTWISE)
BROADCAST_ADD_FUNC(broadcast_a2b, INNER_LOOP_BROADCAST_A2B)
BROADCAST_ADD_FUNC(broadcast_b2a, INNER_LOOP_BROADCAST_B2A)

// void avx2_add_scalar_vector(float scalar, const float *arr, float *result, int len)
// {
//     // 将标量广播到 AVX 寄存器（所有 8 个位置）
//     __m256 scalar_vec = _mm256_set1_ps(scalar);

//     // 主循环：每次处理 8 个 float（AVX 的 256 位寄存器）
//     int i = 0;
//     for (; i + 8 <= len; i += 8)
//     {
//         // 加载 8 个 float（假设内存 32 字节对齐，否则使用 _mm256_loadu_ps）
//         __m256 vec = _mm256_load_ps(arr + i);

//         // 执行向量加法：标量 + 数组元素
//         __m256 sum = _mm256_add_ps(scalar_vec, vec);

//         // 存储结果（同样假设内存对齐）
//         _mm256_store_ps(result + i, sum);
//     }

//     // 处理剩余元素（不足 8 个的部分）
//     for (; i < len; ++i)
//     {
//         result[i] = scalar + arr[i];
//     }
// }
void broadcast_add_3d(float *a_data, int a_shape[3], int a_stride[3],
                       float *b_data, int b_shape[3], int b_stride[3],
                       float *c_data, int c_shape[3], int c_stride[3])
{
    if (a_shape[2] == b_shape[2])
    {
        broadcast_add_eltwise(a_data, a_shape, a_stride, b_data, b_shape, b_stride, c_data, c_shape, c_stride);
    }
    else if (a_shape[2] == 1)
    {
        broadcast_add_broadcast_a2b(a_data, a_shape, a_stride, b_data, b_shape, b_stride, c_data, c_shape, c_stride);
    }
    else if (b_shape[2] == 1)
    {
        broadcast_add_broadcast_b2a(a_data, a_shape, a_stride, b_data, b_shape, b_stride, c_data, c_shape, c_stride);
    }
    else
    {
        broadcast_add_universal(a_data, a_shape, a_stride, b_data, b_shape, b_stride, c_data, c_shape, c_stride);
    }
}